package com.example.ai.service;

import com.example.ai.entity.Document;
import com.example.ai.entity.QARecord;
import com.example.ai.repository.QARecordRepository;
import lombok.RequiredArgsConstructor;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpMethod;
import org.springframework.web.client.RestTemplate;
import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.time.LocalDateTime;
import com.fasterxml.jackson.databind.JsonNode;
import java.lang.StringBuilder;
import com.fasterxml.jackson.core.JsonProcessingException;
import org.springframework.scheduling.annotation.Async;
import lombok.extern.slf4j.Slf4j;

import java.util.List;
import java.util.Map;
import java.util.Arrays;
import java.util.HashMap;

@Service
@Slf4j
@RequiredArgsConstructor
public class QAService {
    @Value("${aliyun.dashscope.api-key}")
    private String apiKey;

    @Value("${deepseek.api-key}")
    private String apiKeyDeepSeek;

    private final QARecordRepository qaRecordRepository;
    private final DocumentService documentService;
    private final DeepseekService deepseekService;
    private final QwenService qwenService;
    private final RestTemplate restTemplate;

    public QARecord askQuestion(Long documentId, String question, String model,Long userId) {
        // 获取文档内容
        Document document = documentService.getDocument(documentId);
        // 构建提示词
        String prompt = buildPrompt(document.getContent(), question);

        // 根据选择的模型调用不同的服务
        String aiResponse;
        if ("deepseek".equals(model)) {
            aiResponse = deepseekService.chat(prompt);
        } else {
            aiResponse = qwenService.chat(prompt);
        }

        // 保存问答记录
        QARecord record = new QARecord();
        record.setDocumentId(documentId);
        record.setQuestion(question);
        record.setAnswer(aiResponse);
        record.setUserId(userId);

        return qaRecordRepository.save(record);
    }

    private String buildPrompt(String documentContent, String question) {
        return String.format("""
            请基于以下文档内容回答问题：
            
            文档内容：
            %s
            
            问题：
            %s
            
            请提供准确、简洁的回答，并尽可能引用文档中的相关内容。
            """, documentContent, question);
    }

    @Async
    public void askStreamQuestion(Long documentId, String question, SseEmitter emitter) {
        log.info("Starting stream question: {}", question);
        try {
            QARecord qaRecord = new QARecord();
            qaRecord.setDocumentId(documentId);
            qaRecord.setQuestion(question);
            qaRecord.setAnswer("");
            qaRecord.setCreateTime(LocalDateTime.now());
            qaRecordRepository.save(qaRecord);

            Map<String, Object> parameters = new HashMap<>();
            parameters.put("stream", true);
            parameters.put("incremental_output", true);

            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "qwen1.5-110b-chat");
            requestBody.put("input", Map.of(
                    "messages", Arrays.asList(
                            Map.of("role", "user", "content", question)
                    )
            ));
            requestBody.put("parameters", parameters);

            StringBuilder fullAnswer = new StringBuilder();

            HttpHeaders headers = new HttpHeaders();
            headers.set("Authorization", "Bearer " + apiKey);
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("X-DashScope-SSE", "enable");

            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

            log.info("Sending request to Qwen API with body: {}", requestBody);

            restTemplate.execute(
                    "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation",
                    HttpMethod.POST,
                    req -> {
                        req.getHeaders().putAll(headers);
                        new ObjectMapper().writeValue(req.getBody(), requestBody);
                    },
                    response -> {
                        try (BufferedReader reader = new BufferedReader(
                                new InputStreamReader(response.getBody()))) {
                            String line;
                            while ((line = reader.readLine()) != null) {
                                log.debug("Raw line: {}", line);

                                if (line.trim().isEmpty()) {
                                    continue;
                                }

                                try {
                                    // 处理 SSE 格式
                                    if (line.startsWith("data:")) {
                                        String jsonData = line.substring(5).trim();
                                        ObjectMapper mapper = new ObjectMapper();
                                        JsonNode node = mapper.readTree(jsonData);
                                        
                                        // 获取文本内容
                                        String content = node.path("output").path("text").asText();
                                        String finishReason = node.path("output").path("finish_reason").asText();
                                        
                                        if (!content.isEmpty()) {
                                            // 直接发送当前文本片段，不做增量处理
                                            log.debug("Sending content: {}", content);
                                            emitter.send(content, MediaType.TEXT_EVENT_STREAM);
                                            fullAnswer.append(content);
                                        }
                                        
                                        if ("stop".equals(finishReason)) {
                                            log.info("Received stop signal");
                                            qaRecord.setAnswer(fullAnswer.toString());
                                            qaRecordRepository.save(qaRecord);
                                            emitter.complete();
                                            break;
                                        }
                                    }
                                } catch (JsonProcessingException e) {
                                    log.warn("Failed to parse JSON: {}", line, e);
                                }
                            }
                        } catch (Exception e) {
                            log.error("Error reading response", e);
                            emitter.completeWithError(e);
                        }
                        return null;
                    }
            );

        } catch (Exception e) {
            log.error("Error in askStreamQuestion", e);
            emitter.completeWithError(e);
        }
    }

    @Async
    public void askStreamQuestionForDeepSeek(Long documentId, String question, SseEmitter emitter) {
        log.info("Starting stream question: {}", question);
        try {
            QARecord qaRecord = new QARecord();
            qaRecord.setDocumentId(documentId);
            qaRecord.setQuestion(question);
            qaRecord.setAnswer("");
            qaRecord.setCreateTime(LocalDateTime.now());
            qaRecordRepository.save(qaRecord);

            Map<String, Object> requestBody = new HashMap<>();
            requestBody.put("model", "deepseek-chat");
            requestBody.put("messages", Arrays.asList(
                            Map.of("role", "user", "content", question)
                    ));

            requestBody.put("stream", true);

            StringBuilder fullAnswer = new StringBuilder();

            HttpHeaders headers = new HttpHeaders();
            headers.set("Authorization", "Bearer " + apiKeyDeepSeek);
            headers.setContentType(MediaType.APPLICATION_JSON);
            headers.set("X-DashScope-SSE", "enable");

            HttpEntity<Map<String, Object>> entity = new HttpEntity<>(requestBody, headers);

            restTemplate.execute(
                    "https://api.deepseek.com/chat/completions",
                    HttpMethod.POST,
                    req -> {
                        req.getHeaders().putAll(headers);
                        new ObjectMapper().writeValue(req.getBody(), requestBody);
                    },
                    response -> {
                        try (BufferedReader reader = new BufferedReader(
                                new InputStreamReader(response.getBody()))) {
                            String line;
                            while ((line = reader.readLine()) != null) {
                                log.debug("Raw line: {}", line);

                                if (line.trim().isEmpty()) {
                                    continue;
                                }

                                try {
                                    // 处理 SSE 格式
                                    if (line.startsWith("data:")) {
                                        String jsonData = line.substring(5).trim();
                                        ObjectMapper mapper = new ObjectMapper();
                                        JsonNode node = mapper.readTree(jsonData);

                                        // 获取文本内容
                                        String content = node.get("choices").get(0).get("delta").get("content").asText() + " ";
                                        String finishReason = node.get("choices").get(0).get("finish_reason").asText();

                                        if (!content.isEmpty()) {
                                            // 直接发送当前文本片段，不做增量处理
                                            log.debug("Sending content: {}", content);
                                            emitter.send(content, MediaType.TEXT_EVENT_STREAM);
                                            fullAnswer.append(content);
                                        }

                                        if ("stop".equals(finishReason)) {
                                            log.info("Received stop signal");
                                            qaRecord.setAnswer(fullAnswer.toString());
                                            qaRecordRepository.save(qaRecord);
                                            emitter.complete();
                                            break;
                                        }
                                    }
                                } catch (JsonProcessingException e) {
                                    log.warn("Failed to parse JSON: {}", line, e);
                                }
                            }
                        } catch (Exception e) {
                            log.error("Error reading response", e);
                            emitter.completeWithError(e);
                        }
                        return null;
                    }
            );

        } catch (Exception e) {
            log.error("Error in askStreamQuestion", e);
            emitter.completeWithError(e);
        }
    }

    public List<QARecord> getDocumentQAHistory(Long documentId, Long userId) {
        return qaRecordRepository.findByUserIdAndDocumentIdOrderByCreateTimeDesc(userId, documentId);
    }

    public List<QARecord> getUserQAHistory(Long userId) {
        return qaRecordRepository.findByUserIdOrderByCreateTimeDesc(userId);
    }

} 